import networkx as nx
from typing import Dict, Any, List
from app.services.dhgam.thgad_service import TemporalHeterogeneousGraphAnomalyDetection
from app.schemas.dhgam.thgad_schema import (
    InputParams, OutputParams, AlgorithmResponse, 
    THGADOutputParams, AnomalyResult, TaskStatus
)
from app.utils.logger import logger

class THGADService:
    def __init__(self):
        """初始化THGAD服务"""
        pass
    
    async def run_thgad(self, input_params: Dict[str, Any], task_id: str) -> AlgorithmResponse:
        """
        运行THGAD算法
        
        Args:
            input_params: 算法输入参数
            task_id: 任务ID
            
        Returns:
            AlgorithmResponse: 算法执行结果
        """
        try:
            # 提取算法参数
            thgad_params = input_params.get("thgad_params", {})
            
            # 从输入参数中构建图序列
            graph_sequence_data = thgad_params.get("graph_sequence", [])
            graph_sequence = self._build_graph_sequence(graph_sequence_data)
            
            # 设置算法参数
            window_size = thgad_params.get("window_size", 5)
            contamination = thgad_params.get("contamination", 0.1)
            
            # 初始化并运行THGAD算法
            thgad = TemporalHeterogeneousGraphAnomalyDetection(
                window_size=window_size,
                contamination=contamination
            )
            thgad.fit(graph_sequence)
            anomalies_dict = thgad.detect(graph_sequence)
            
            # 构建输出结果
            anomalies = []
            for node, score in anomalies_dict.items():
                node_type = graph_sequence[-1].nodes[node]['type']
                anomalies.append(AnomalyResult(
                    node_id=node,
                    anomaly_score=float(score),
                    node_type=node_type
                ))
            
            # 创建算法输出参数
            parameters = {
                "window_size": window_size,
                "contamination": contamination
            }
            thgad_results = THGADOutputParams(
                anomalies=anomalies,
                parameters=parameters
            )
            
            # 构建响应对象
            response = AlgorithmResponse(
                task_id=task_id,
                task_status=TaskStatus.COMPLETED,
                task_progress=100,
                output_params=OutputParams(thgad_results=thgad_results),
                metrics=self._generate_metrics(anomalies)
            )
            
            return response
            
        except Exception as e:
            logger.error(f"THGAD算法执行失败: {str(e)}")
            return AlgorithmResponse(
                task_id=task_id,
                task_status=TaskStatus.FAILED,
                error_message=f"THGAD算法执行失败: {str(e)}",
                output_params=OutputParams()
            )
    
    def _build_graph_sequence(self, graph_sequence_data: List[Dict[str, Any]]) -> List[nx.Graph]:
        """
        从输入数据构建NetworkX图序列
        
        Args:
            graph_sequence_data: 图序列数据
            
        Returns:
            List[nx.Graph]: NetworkX图对象列表
        """
        graph_sequence = []
        for graph_data in graph_sequence_data:
            G = nx.Graph()
            
            # 添加节点
            nodes = graph_data.get("nodes", {})
            for node_type, node_list in nodes.items():
                for node_info in node_list:
                    node_id = node_info.get("id")
                    # 确保节点属性是字典，并添加节点类型
                    node_attrs = {k: v for k, v in node_info.items() if k != "id"}
                    node_attrs["type"] = node_type
                    G.add_node(node_id, **node_attrs)
            
            # 添加边
            for edge in graph_data.get("edges", []):
                source = edge.get("source")
                target = edge.get("target")
                attrs = {k: v for k, v in edge.items() if k not in ["source", "target"]}
                G.add_edge(source, target, **attrs)
            
            graph_sequence.append(G)
        
        return graph_sequence
    
    def _generate_metrics(self, anomalies: List[AnomalyResult]) -> List[Any]:
        """
        生成可视化图表
        
        Args:
            anomalies: 检测到的异常列表
            
        Returns:
            List[Chart]: 图表列表
        """
        # 这里可以生成各种可视化图表，如异常分布图等
        # 由于Chart类的具体实现不详，这里返回空列表
        return [] 